"""
简单游戏演示
提供一个基础的游戏循环和可视化界面，用于快速测试和演示

作者：wjttdbx
版本：3.0 - 适配重构后的代码
"""
import sys
import pygame

sys.path.insert(0, sys.path[0]+"/../")
from core.config import GameConfig
from agents import white_agent
from core.CoreEnvironment import CoreEnvironment
from agents.fsm_defense_agent import create_fsm_defense_agent
from agents.simple_black_agent import SimpleBlackAgent
from agents.base_agent import AgentObservation
from agents.action_executor import ActionExecutor
from visualization import GameVisualizer


class SimpleGameDemo:
    """
    简单游戏演示类
    
    提供一个基础的游戏循环和可视化界面
    """
    def __init__(self, width=1280, height=720, use_fsm_agent=False, black_agent_type="simple"):
        """初始化游戏"""
        pygame.init()
        self.width = width
        self.height = height
        self.screen = pygame.display.set_mode((width, height))
        pygame.display.set_caption("SH-15 博弈决策模型演示")
        self.core_env = CoreEnvironment()
        


        self.white_agent = create_fsm_defense_agent()
        self.agent_type = "FSM防御智能体"

            
        self.black_agent = SimpleBlackAgent()
        self.action_executor = ActionExecutor()
        
        # 启用调试模式以查看详细的动作执行信息
        self.action_executor.set_debug_mode(GameConfig.debug_mode)

        self.visualization = GameVisualizer(self.screen)
        
        self.running = False
        self.total_steps = 0

    def run_game_loop(self):
        """运行主游戏循环"""
        self.running = True
        clock = pygame.time.Clock()
        
        # 打印策略说明
        print(f"=== SH-15 {self.agent_type}  ===")
        print("策略说明:")
        

        # 白方策略说明
        if "FSM" in self.agent_type:
            print("- 白方: 使用FSM状态机控制，智能部署防线和无人机巡航")
            print("- 无人机: 按指定路径巡航，在关键点进行圆形扫描")
        else:
            print("- 白方: 在x=80000线上均匀部署，无人机水平巡航至x=25000后返回，锁定探测到的目标")
        print("==============================\n")
        
        # 重置环境和智能体
        self.core_env.reset()
        self.white_agent.reset()
        self.black_agent.reset()
        
        print("游戏初始化完成")
        print(f"白方无人艇数量: {len(self.core_env.white_boats)}")
        print(f"黑方无人艇数量: {len(self.core_env.black_boats)}")
        print("开始游戏循环...\n")

        print("===白方初始为部署阶段===")
        
        while self.running:
            # 控制渲染帧率，但不影响游戏时间步长
            clock.tick(60)
            
            # 使用固定的游戏时间步长，并结合可视化器的速度倍数
            dt = GameConfig.TIME_STEP * self.visualization.speed_multiplier * 0.1   # 人为调整下速度

            self._handle_input()
            self._update_game_state(dt)
            self._render()
            self._check_game_over()
            
            self.total_steps += 1

    def _update_game_state(self, dt: float):
        """更新游戏状态"""
        # 1. 更新探测系统，让AI获取最新信息
        self.core_env.update_detection_system()
        
        # 2. AI决策
        # 创建观测对象
        # 仅传入“当前帧被探测到”的黑方目标，避免越权获取真实坐标
        detected_ids = set()
        for boat in self.core_env.white_boats:
            for t in boat.detected_targets:
                if t.active:
                    detected_ids.add(t.id)
        for drone in self.core_env.all_white_drones:
            for t in drone.detected_targets:
                if t.active:
                    detected_ids.add(t.id)

        detected_black_boats = [
            b for b in self.core_env.black_boats
            if b.active and b.id in detected_ids
        ]

        observation = AgentObservation(
            white_boats=self.core_env.white_boats,
            all_drones=self.core_env.all_white_drones,
            black_boats=detected_black_boats,
            current_time=self.core_env.current_time
        )       
        agent_action = self.white_agent.get_actions(observation)
        if GameConfig.debug_mode:
            print(f"\n第{self.total_steps}步:")
            print(f"白方无人艇动作: {agent_action.boat_actions}")
            print(f"白方无人机动作: {agent_action.drone_actions}")
            print(f"当前时间: {self.core_env.current_time:.2f} 秒")


        # if self.total_steps % 100 == 0:
        #     print(f"历史探测信息: {self.white_agent.decision_history}")
        
        # 3. 执行白方动作
        self.action_executor.execute_actions(
            agent_action,
            self.core_env.white_boats,
            self.core_env.all_white_drones,
            self.core_env,
            self.core_env.current_time,
            dt
        )
        
        # 打印白方无人艇和无人机的真实速度和航向
        if GameConfig.debug_mode:
            print(f"第{self.total_steps}步 -")
            for i, boat in enumerate(self.core_env.white_boats):
                print(f"  白船{i} 速度: {boat.speed:.2f} m/s, 航向: {boat.heading:.2f} 弧度")
            for i, drone in enumerate(self.core_env.all_white_drones):
                print(f"  白机{i} 速度: {drone.speed:.2f} m/s, 航向: {drone.heading:.2f} 弧度")
        # 4. 更新黑方单位
        self.black_agent.update_all_black_boats(
            self.core_env.black_boats,
            self.core_env.white_boats,
            self.core_env.current_time,
            dt
        )
        # print(self.white_agent.predict_black_positions(None))
        

        # 5. 更新核心环境（物理、碰撞等）
        self.core_env.update(dt)

    def _handle_input(self):
        """处理用户输入"""
        for event in pygame.event.get():
            if not self.visualization.handle_event(event):
                self.running = False

    def _render(self):
        """渲染游戏画面"""
        game_state_info = self.core_env.get_game_state()
        
        # 获取状态机信息（统一构造，至少包含 predicted_black_targets）
        # 提取每帧预测位置映射：{black_id: (px, py)}，优先 predict_pos，退回 position
        predicted_black_targets = {}
        try:
            detected_dict = getattr(self.white_agent, 'detected_targets', {}) or {}
            for unit_id, data in detected_dict.items():
                pos = data.get('predict_pos', data.get('position'))
                if pos is None or len(pos) < 2:
                    continue
                try:
                    key = int(unit_id)
                except Exception:
                    key = str(unit_id)
                predicted_black_targets[key] = (float(pos[0]), float(pos[1]))
        except Exception:
            predicted_black_targets = {}

        extras = {}
        if hasattr(self.white_agent, 'global_state') and hasattr(self.white_agent, 'boat_states'):
            task_manager = getattr(self.white_agent, 'task_manager', None)
            global_state = getattr(self.white_agent, 'global_state', None)
            extras = {
                'global_state': global_state.name if global_state else 'UNKNOWN',
                'state_timer': getattr(self.white_agent, 'state_timer', 0),
                'boat_states': {bid: state.name for bid, state in getattr(self.white_agent, 'boat_states', {}).items()},
                'drone_states': {did: state.name for did, state in getattr(self.white_agent, 'drone_states', {}).items()},
                'boat_targets': getattr(self.white_agent, 'boat_targets', {}).copy(),
                'task_assignment_ready': (task_manager is not None and 
                                          hasattr(task_manager, 'is_ready') and
                                          task_manager.is_ready()),
            }

        # 在交给可视化前，按当前活跃黑方过滤 predicted_black_targets，避免显示已击沉目标
        try:
            active_black_ids = {b.id for b in self.core_env.black_boats if b.active}
            filtered = {}
            for uid, pos in predicted_black_targets.items():
                try:
                    uid_int = int(uid)
                except Exception:
                    uid_int = None
                if (uid_int is not None and uid_int in active_black_ids) or (uid in active_black_ids) or (str(uid) in {str(x) for x in active_black_ids}):
                    filtered[uid] = pos
            predicted_black_targets = filtered
        except Exception:
            pass

        fsm_info = {'predicted_black_targets': predicted_black_targets, **extras}
        
        self.visualization.render_frame(
            white_boats=self.core_env.white_boats,
            white_drones=self.core_env.all_white_drones,
            black_boats=self.core_env.black_boats,
            current_time=self.core_env.current_time,
            game_info={ # 传递额外信息
                'collisions': self.core_env.collision_count,
                'intercepts': self.core_env.intercepted_black_boats,
                'fsm_info': fsm_info  # 新增：状态机信息
            }
        )
        pygame.display.flip()

    def _check_game_over(self):
        """检查游戏是否结束"""
        game_over, winner = self.core_env.is_game_over()
        if game_over:
            print(f"\n游戏结束: {winner}")
            self._show_game_over_screen(winner)
            self.running = False

    def _show_game_over_screen(self, winner: str):
        """显示游戏结束画面"""
        game_stats = self.core_env.get_game_state()
        
        if winner == "white_wins":
            result_text = "白方胜利 (White Wins)"
        elif winner == "black_wins":
            result_text = "黑方胜利 (Black Wins)"
        else:
            result_text = "平局 (Draw)"

        # 控制台打印关键统计信息
        crossed_black = game_stats.get('crossed_black_boats', 0)
        collision_count = game_stats.get('collision_count', 0)
        white_lock_count = game_stats.get('white_lock_count', 0)
        black_lock_count = game_stats.get('black_lock_count', 0)
        first_break_time = game_stats.get('first_breakthrough_time', None)
        if first_break_time is None:
            intercept_time_text = "100小时"
        else:
            total_seconds = int(first_break_time)
            minutes = total_seconds // 60
            seconds = total_seconds % 60
            intercept_time_text = f"{minutes}分{seconds:02d}秒"

        print("\n=== 关键统计（控制台） ===")
        print(f"黑方突防成功数量: {crossed_black}")
        print(f"白方阻击时间: {intercept_time_text}")
        print(f"白方被锁定次数: {black_lock_count}")
        print(f"黑方被锁定次数: {white_lock_count}")
        print(f"白方发生碰撞次数: {collision_count}")
        print("======================\n")

        self._render()

        self.visualization.draw_game_over_overlay(
            result=result_text,
            game_stats={
                'steps': self.total_steps,
                'crossed_black_boats': game_stats.get('crossed_black_boats', 0),
                'eliminated_black_boats': game_stats.get('eliminated_black_boats', 0),
                'intercepted_black_boats': game_stats.get('intercepted_black_boats', 0),
                'collision_count': game_stats.get('collision_count', 0),
                'white_lock_count': game_stats.get('white_lock_count', 0),
                'black_lock_count': game_stats.get('black_lock_count', 0),
                'first_breakthrough_time': game_stats.get('first_breakthrough_time', None),
            },
            current_time=self.core_env.current_time
        )
        
        waiting_for_exit = True
        while waiting_for_exit:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    waiting_for_exit = False
                if event.type == pygame.KEYDOWN and (event.key == pygame.K_RETURN or event.key == pygame.K_ESCAPE):
                    waiting_for_exit = False
    
    def cleanup(self):
        """清理资源"""
        pygame.quit()
        print("资源清理完成")


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description='SH-15博弈决策模型演示')
    parser.add_argument('--agent', choices=['simple', 'fsm'], default='simple',
                        help='选择白方智能体类型: simple(简单防线) 或 fsm(FSM状态机)')
    parser.add_argument('--black', choices=['simple', 'coordinated', 'swarm'], default='simple',
                        help='选择黑方智能体类型: simple(简单直行), coordinated(协同作战), swarm(集群围猎)')
    parser.add_argument('--debug', action='store_true', help='启用调试模式')
    
    args = parser.parse_args()
    
    # 设置调试模式
    if args.debug:
        GameConfig.debug_mode = True
    
    use_fsm = (args.agent == 'fsm')
    demo = SimpleGameDemo(use_fsm_agent=use_fsm, black_agent_type=args.black)
    
    print(f"启动演示 - 白方: {demo.agent_type}")
    print("按 ESC 或关闭窗口退出")
    print("=" * 50)
    
    try:
        demo.run_game_loop()
    except (KeyboardInterrupt, SystemExit):
        print("\n用户中断游戏")
    except Exception as error:
        print(f"\n游戏出现错误: {error}")
        import traceback
        traceback.print_exc()
    finally:
        demo.cleanup()
        